/**
 * 
 */
package cn.edu.dutir.test.unit.terrier;

import gnu.trove.TIntArrayList;
import gnu.trove.TIntDoubleHashMap;
import gnu.trove.TIntIntHashMap;


import java.util.Arrays;

import org.apache.log4j.Logger;

import uk.ac.gla.terrier.matching.MatchingQueryTerms;
import uk.ac.gla.terrier.matching.ResultSet;
import uk.ac.gla.terrier.matching.models.queryexpansion.QueryExpansionModel;
import uk.ac.gla.terrier.querying.Manager;
import uk.ac.gla.terrier.querying.PostProcess;
import uk.ac.gla.terrier.querying.Request;
import uk.ac.gla.terrier.querying.SearchRequest;
import uk.ac.gla.terrier.structures.CollectionStatistics;
import uk.ac.gla.terrier.structures.DirectIndex;
import uk.ac.gla.terrier.structures.DocumentIndex;
import uk.ac.gla.terrier.structures.Index;
import uk.ac.gla.terrier.structures.InvertedIndex;
import uk.ac.gla.terrier.structures.Lexicon;
import uk.ac.gla.terrier.structures.LexiconEntry;
import uk.ac.gla.terrier.utility.ApplicationSetup;
import uk.ac.gla.terrier.utility.HeapSort;

/**
 * Referring to the algorithm, pleas see: I2R at NTCIR5
 * @author YeZheng
 *
 */
public class I2RRerankPostProcess implements PostProcess {

	protected static Logger logger = Logger.getRootLogger();
	/** The document index used for retrieval. */
	protected DocumentIndex documentIndex;
	/** The inverted index used for retrieval. */
	protected InvertedIndex invertedIndex;
	/** An instance of Lexicon class. */
	protected Lexicon lexicon;
	/** The direct index used for retrieval. */
	protected DirectIndex directIndex;
	/** The statistics of the index */
	protected CollectionStatistics collStats;
	/** The query expansion model used. */
	protected QueryExpansionModel QEModel;
	protected Request request;
	protected String queryID;
	
	
	double i2rScorces[];
	double scores[];
	private short[] occurences;
	private int[] docids;
	private int set_size;
	
	static int reRankNum = Integer.parseInt(ApplicationSetup.getProperty("I2RRerank.rerankNum", "100"));
	/* (non-Javadoc)
	 * @see uk.ac.gla.terrier.querying.PostProcess#getInfo()
	 */
	public String getInfo() {
		return "I2RRerankPostProcess_" + "rNum=" + reRankNum;
	}

	public void getI2RScores(Manager manager, SearchRequest q){
	   	Index index = manager.getIndex();
		documentIndex = index.getDocumentIndex();
		invertedIndex = index.getInvertedIndex();
		lexicon = index.getLexicon();
		collStats = index.getCollectionStatistics(); 
		int num_doc = collStats.getNumberOfDocuments();
		directIndex = index.getDirectIndex();
		this.request = (Request)q;
		this.queryID = this.request.getQueryID();

		ResultSet resultSet = q.getResultSet();
		set_size = resultSet.getExactResultSize();
		docids = resultSet.getDocids();
		scores = resultSet.getScores();
		occurences = resultSet.getOccurrences();
		
		
		this.i2rScorces = new double[set_size];
		Arrays.fill(this.i2rScorces, 0d);
		
		TIntIntHashMap q_dfMap = new TIntIntHashMap();
		TIntIntHashMap qlenMap = new TIntIntHashMap();
		MatchingQueryTerms mqts = this.request.getMatchingQueryTerms();
		String[] terms = mqts.getTerms();
		int terms_size = terms.length;
		for (int i = 0; i < terms.length; i++) {			
			LexiconEntry lEntry = this.lexicon.getLexiconEntry(terms[i]);
			if(lEntry == null){
				continue;
			}
			q_dfMap.put(lEntry.termId,lEntry.n_t);
			qlenMap.put(lEntry.termId,terms[i].length());
		}
		
		TIntDoubleHashMap termWeights = new TIntDoubleHashMap();
		TIntArrayList docTermsLists[] = new TIntArrayList[set_size];
		for (int i = 0; i < set_size; i++) {
			docTermsLists[i] = new TIntArrayList();
			int t_tfs[][] = directIndex.getTerms(docids[i]);
			for (int k = 0; k < t_tfs[0].length; k++) {
				int termid = t_tfs[0][k];
				if(q_dfMap.contains(termid)){
					docTermsLists[i].add(termid);
					double tvalue = 1d + 1/Math.sqrt((double)(1d + i));
					termWeights.adjustOrPutValue(termid, tvalue, tvalue);
				}
			}
		}
		
		for (int i = 0; i < set_size; i++) {
			int len = docTermsLists[i].size();
			for(int j=0; j < len; j++){
				int termid = docTermsLists[i].get(j);
				this.i2rScorces[i] += Math.sqrt(termWeights.get(termid)*num_doc/(set_size*q_dfMap.get(termid)))
						*Math.sqrt(qlenMap.get(termid));
			}
		}
	}
	
	/* (non-Javadoc)
	 * @see uk.ac.gla.terrier.querying.PostProcess#process(uk.ac.gla.terrier.querying.Manager, uk.ac.gla.terrier.querying.SearchRequest)
	 */
	public void process(Manager manager, SearchRequest q) {
		logger.info("Starting I2RRerankPostProcess post-processing.");
		getI2RScores(manager, q);
		for(int i=0; i < this.set_size; i++){
			double tv = this.scores[i];
			this.scores[i] = this.scores[i]*(1 + this.i2rScorces[i]);
			logger.info(tv + " , " + this.i2rScorces[i] + " , " + this.scores[i]);
		}
		int rnum = Math.min(reRankNum, set_size);
		HeapSort.descendingHeapSort(scores, docids, occurences, rnum);
	}

	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		// TODO Auto-generated method stub

	}
}
